// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/exec/arrow_stream.h>
#include <pollux/vector/arrow/bridge.h>

namespace kumo::pollux::exec {
    ArrowStream::ArrowStream(
        int32_t operatorId,
        DriverCtx *driverCtx,
        const std::shared_ptr<const core::ArrowStreamNode> &arrowStreamNode)
        : SourceOperator(
            driverCtx,
            arrowStreamNode->outputType(),
            operatorId,
            arrowStreamNode->id(),
            "ArrowStream") {
        arrowStream_ = arrowStreamNode->arrowStream();
    }

    ArrowStream::~ArrowStream() {
        close();
    }

    RowVectorPtr ArrowStream::getOutput() {
        // Get Arrow array.
        struct ArrowArray arrowArray;
        if (arrowStream_->get_next(arrowStream_.get(), &arrowArray)) {
            if (arrowArray.release) {
                arrowArray.release(&arrowArray);
            }
            POLLUX_FAIL(
                "Failed to call get_next on ArrowStream: {}", std::string(getError()));
        }
        if (arrowArray.release == nullptr) {
            // End of Stream.
            finished_ = true;
            return nullptr;
        }

        // Get Arrow schema.
        struct ArrowSchema arrowSchema;
        if (arrowStream_->get_schema(arrowStream_.get(), &arrowSchema)) {
            if (arrowSchema.release) {
                arrowSchema.release(&arrowSchema);
            }
            if (arrowArray.release) {
                arrowArray.release(&arrowArray);
            }
            POLLUX_FAIL(
                "Failed to call get_schema on ArrowStream: {}",
                std::string(getError()));
        }

        // Convert Arrow Array into RowVector and return.
        return std::dynamic_pointer_cast<RowVector>(
            import_from_arrow_as_owner(arrowSchema, arrowArray, pool()));
    }

    bool ArrowStream::isFinished() {
        return finished_;
    }

    const char *ArrowStream::getError() const {
        const char *lastError = arrowStream_->get_last_error(arrowStream_.get());
        POLLUX_CHECK_NOT_NULL(lastError);
        return lastError;
    }

    void ArrowStream::close() {
        if (arrowStream_->release) {
            arrowStream_->release(arrowStream_.get());
        }
        SourceOperator::close();
    }
} // namespace kumo::pollux::exec
